#include "Irp.h"
#include "Util.h"

NTSTATUS IrpGeneral(IN PDEVICE_OBJECT DeviceObject, IN PIRP Irp)
{
    PDEVICE_EXTENSION lpDeviceExtension = (PDEVICE_EXTENSION)(DeviceObject->DeviceExtension);
    NTSTATUS retStatus = STATUS_SUCCESS;

    switch (lpDeviceExtension->DeviceType)
    {
    case CONTROL:
    {
        CompleteRequest(Irp, retStatus, 0);
        break;
    }
    case FILTER:
    {
        IoSkipCurrentIrpStackLocation(Irp);
        retStatus = IoCallDriver(lpDeviceExtension->AttachedToDeviceObject, Irp);
        break;
    }
    default:
        break;
    }

    return retStatus;
}

static NTSTATUS ReadCompletionRoutine(IN PDEVICE_OBJECT DeviceObject, IN PIRP Irp, IN PVOID Context OPTIONAL)
{
    UNREFERENCED_PARAMETER(Context);
    UNREFERENCED_PARAMETER(DeviceObject);

    PKEYBOARD_INPUT_DATA lpKbInputData;
    ULONG_PTR uCnt;
    char strCode[20];

    if (Irp->PendingReturned)
        IoMarkIrpPending(Irp);

    if (Irp->IoStatus.Information)
    {
        lpKbInputData = Irp->AssociatedIrp.SystemBuffer;
        uCnt = Irp->IoStatus.Information / sizeof(KEYBOARD_INPUT_DATA);
        //if (lpKbInputData->Flags & KEY_BREAK)
        //{
        //    DBGPRINT("%X Up\n", lpKbInputData->MakeCode);
        //}
        //else
        //{
        //    DBGPRINT("%X Down\n", lpKbInputData->MakeCode);
        //}

        for (unsigned int i = 0; i < uCnt; ++i)
        {
            RtlZeroMemory(strCode, sizeof(strCode));
            MakeCodeToASCII(lpKbInputData[i].MakeCode, lpKbInputData[i].Flags, strCode);
            DBGPRINT("%s\n", strCode);
        }
    }

    IoReleaseRemoveLock(&KfGlobals.RemoveLock, TAG_USE_REMOVE_LOCK);
    return STATUS_CONTINUE_COMPLETION;
}

NTSTATUS IrpRead(IN PDEVICE_OBJECT DeviceObject, IN PIRP Irp)
{
    PDEVICE_EXTENSION lpDeviceExtension = (PDEVICE_EXTENSION)(DeviceObject->DeviceExtension);
    NTSTATUS retStatus = STATUS_SUCCESS;
    NTSTATUS ntStatus;

    switch (lpDeviceExtension->DeviceType)
    {
    case CONTROL:
    {
        CompleteRequest(Irp, retStatus, 0);
        break;
    }
    case FILTER:
    {
        ntStatus = IoAcquireRemoveLock(&KfGlobals.RemoveLock, TAG_USE_REMOVE_LOCK);
        if (STATUS_SUCCESS != ntStatus) break;

        IoCopyCurrentIrpStackLocationToNext(Irp);
        IoSetCompletionRoutine(Irp, ReadCompletionRoutine, NULL, TRUE, TRUE, TRUE);
        retStatus = IoCallDriver(lpDeviceExtension->AttachedToDeviceObject, Irp);
        break;
    }
    default:
        break;
    }

    return retStatus;
}

static NTSTATUS KbReadCompletionRoutine(IN PDEVICE_OBJECT DeviceObject, IN PIRP Irp, IN PVOID Context OPTIONAL)
{
    PKEYBOARD_INPUT_DATA lpKbInputData;
    ULONG_PTR uCnt;
    char strCode[20];

    if (Irp->PendingReturned)
        IoMarkIrpPending(Irp);

    if (Irp->IoStatus.Information)
    {
        lpKbInputData = Irp->AssociatedIrp.SystemBuffer;
        uCnt = Irp->IoStatus.Information / sizeof(KEYBOARD_INPUT_DATA);
        //if (lpKbInputData->Flags & KEY_BREAK)
        //{
        //    DBGPRINT("%X Up\n", lpKbInputData->MakeCode);
        //}
        //else
        //{
        //    DBGPRINT("%X Down\n", lpKbInputData->MakeCode);
        //}

        for (unsigned int i = 0; i < uCnt; ++i)
        {
            RtlZeroMemory(strCode, sizeof(strCode));
            MakeCodeToASCII(lpKbInputData[i].MakeCode, lpKbInputData[i].Flags, strCode);
            DBGPRINT("%s\n", strCode);
        }
    }

    if (KfGlobals.KbIrpReadCompletionRoutine)
        KfGlobals.KbIrpReadCompletionRoutine(DeviceObject, Irp, Context);

    IoReleaseRemoveLock(&KfGlobals.RemoveLock, TAG_USE_REMOVE_LOCK);
    return STATUS_CONTINUE_COMPLETION;
}

NTSTATUS KbIrpReadHook(IN PDEVICE_OBJECT DeviceObject, IN PIRP Irp)
{
    NTSTATUS ntStatus;
    PIO_STACK_LOCATION lpIoStackLocation = IoGetCurrentIrpStackLocation(Irp);

    ntStatus = IoAcquireRemoveLock(&KfGlobals.RemoveLock, TAG_USE_REMOVE_LOCK);
    if (STATUS_SUCCESS != ntStatus) goto _retIrpReadHook;

    lpIoStackLocation->Control = SL_INVOKE_ON_CANCEL | SL_INVOKE_ON_ERROR | SL_INVOKE_ON_SUCCESS;
    KfGlobals.KbIrpReadCompletionRoutine = (PIO_COMPLETION_ROUTINE)InterlockedExchange64((LONG64*)(&(lpIoStackLocation->CompletionRoutine)),
        (LONG64)KbReadCompletionRoutine);

    if (KfGlobals.KbIrpRead)
        return KfGlobals.KbIrpRead(DeviceObject, Irp);

_retIrpReadHook:
    CompleteRequest(Irp, STATUS_SUCCESS, 0);
    return STATUS_SUCCESS;
}

NTSTATUS IrpPnp(IN PDEVICE_OBJECT DeviceObject, IN PIRP Irp)
{
    PIO_STACK_LOCATION lpIoStackLocation = IoGetCurrentIrpStackLocation(Irp);
    PDEVICE_EXTENSION lpDeviceExtension = (PDEVICE_EXTENSION)DeviceObject->DeviceExtension;

    if (lpIoStackLocation->MinorFunction == IRP_MN_REMOVE_DEVICE)
    {
        IoSkipCurrentIrpStackLocation(Irp);
        IoCallDriver(lpDeviceExtension->AttachedToDeviceObject, Irp);
        IoDetachDevice(DeviceObject);
        IoDeleteDevice(DeviceObject);
        return STATUS_SUCCESS;
    }

    IoSkipCurrentIrpStackLocation(Irp);
    return IoCallDriver(lpDeviceExtension->AttachedToDeviceObject, Irp);
}

NTSTATUS IrpPower(IN PDEVICE_OBJECT DeviceObject, IN PIRP Irp)
{
    PDEVICE_EXTENSION lpDeviceExtension = (PDEVICE_EXTENSION)DeviceObject->DeviceExtension;
    PoStartNextPowerIrp(Irp);
    IoSkipCurrentIrpStackLocation(Irp);
    return PoCallDriver(lpDeviceExtension->AttachedToDeviceObject, Irp);
}

